import numpy as np
import scipy.sparse as sp

from quimb.core import (ensure_qarray)
@ensure_qarray
def _diag_dense(v, k=0, dtype=complex):
	"""Returns a dense, diagonal matrix of given diagonal ``v``, offset ``k`` and type ``dtype``."""
	return np.diag(v, k=k).astype(dtype)

def _diag_sparse(v, k=0, stype="csr", dtype=complex):
	"""Returns a sparse, real diagonal matrix."""
	return sp.diags(v, offsets=k, format=stype, dtype=dtype)

def diag(v, k=0, sparse=False, stype="csr", dtype=complex):
	if sparse:
		return _diag_sparse(v, k=k, stype=stype, dtype=dtype)
	return _diag_dense(v, k=k, dtype=dtype)

